package com.le.enhance.class7;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

// 并查集
public class Code_04_UnionFindSet {
    // 元素类型,
    public static class Node<V> {
        V value;

        public Node(V value) {
            this.value = value;
        }
    }

    public static class UnionSet<V> {
        // 元素val --> node
        public HashMap<V, Node<V>> nodeMap;
        // 当前node的父节点
        public HashMap<Node<V>, Node<V>> fatherMap;
        // node 大小
        public HashMap<Node<V>, Integer> sizeMap;

        // 初始化
        public UnionSet(List<V> list) {
            this.nodeMap = new HashMap<>();
            this.fatherMap = new HashMap<>();
            this.sizeMap = new HashMap<>();
            for (V val : list) {
                Node<V> node = new Node<>(val);
                this.nodeMap.put(val, node); // 通过当前元素新建结点
                this.fatherMap.put(node, node); // 当前结点父节点为自己
                this.sizeMap.put(node, 1); // 集合大小为1
            }
        }

        // 查找当前node的父节点
        // 优化, 把路径上的所有点都指向父节点
        public Node<V> findHead(Node<V> node) {
            Stack<Node<V>> stack = new Stack<>();
            while (node != fatherMap.get(node)) {
                stack.push(node);
                node = fatherMap.get(node);
            }
            while (!stack.isEmpty()) {
                this.fatherMap.put(stack.pop(), node);
            }
            return node;
        }

        // 判断两个元素是否属于同一集合
        public boolean isSampleSet(V a, V b) {
            // 判断元素是否在集合中
            if (nodeMap.containsKey(a) && nodeMap.containsKey(b)) {
                // 判断父节点是否相同
                return findHead(nodeMap.get(a)) == findHead(nodeMap.get(b));
            }
            return false;
        }

        // 合并两个集合
        public void union(V a, V b) {
            if (nodeMap.containsKey(a) && nodeMap.containsKey(b)) {
                // 获取两者的父节点
                Node<V> aF = findHead(nodeMap.get(a));
                Node<V> bF = findHead(nodeMap.get(b));
                if (aF != bF) {
                    // 判断大小, size小 --> size大, 删除size小的
                    Node<V> big = sizeMap.get(aF) >= sizeMap.get(bF) ? aF : bF;
                    Node<V> small = big == aF ? bF : aF;
                    fatherMap.put(small, big);
                    sizeMap.put(big, sizeMap.get(aF) + sizeMap.get(bF));
                    fatherMap.remove(small);
                }
            }
        }
    }
}
